import os
import multiprocessing
import shlex
import shutil

import subprocess
import random

import math
from tasks import pre_pocess
import lmdb  # May require 'pip install lmdb' if lmdb not found
import numpy as np
import logging
from logging import info, error
import scipy.io


FORMAT = '[%(levelname)-5s]%(asctime)-8s %(filename)s:%(lineno)d %(message)s'
DATEFORMAT = '%H:%M:%S'
logging.basicConfig(level=logging.DEBUG, format=FORMAT, datefmt=DATEFORMAT)
logging.debug('start')

home = os.getenv('HOME', '')
caffe_root = home + '/caffe'
tools_root = caffe_root + '/build/tools'
import sys

sys.path.insert(0, caffe_root + '/python')
import caffe

num_thread = 50
# num_thread = 8




f = open(home + '/data/grade/trainLabels.csv')
home = os.getenv('HOME', '')
o_root = 'out'
dbroot = home + '/data/grade/train'
lines = f.readlines()[1:]

print('listing directory ...')
exist_file = os.listdir(dbroot)
random.shuffle(exist_file)
# exist_file = exist_file[:20]
if not os.path.exists(o_root):
    os.mkdir(o_root)

print('parsing file ...')
file_level = {}
for line in lines:
    line = line[:-1]
    strs = line.split(',')
    fname = strs[0] + '.jpeg'
    level = int(strs[1])
    file_level[fname] = level

train_number = math.floor(len(exist_file) * 0.8)
for phase, files in [
    ('train', exist_file[:train_number]),
    ('test', exist_file[train_number:])
]:
    print('running phase ', phase, '....')
    # genrate list file
    list_file_name = o_root + '/filelist.' + phase
    img_root = o_root + '/'
    tlist = {}
    print('generating list file')
    with open(list_file_name, 'w') as f:
        index = 0
        for fname in files:
            oname = os.path.splitext(fname)[0] + '.png'
            if os.path.exists(o_root + '/' + oname):
                index += 1
                continue
            tlist[fname] = pre_pocess.delay(dbroot + '/' + fname, o_root + '/' + oname)
            while len(tlist) > num_thread:
                for fn in tlist:
                    try:
                        if tlist[fn].ready():
                            v = tlist[fn].get()
                            if v != '':
                                f.write('%s %d' % (v, 0) + '\n')
                                index+=1
                            else:
                                print('error', fn)
                                files.remove(fn)
                            tlist.pop(fn)
                            break
                    except Exception as e:
                        error(e)
            if index % 100 == 0:
                print('finished', index,sep=' ', end='\r')
        for fn in tlist:
            v = tlist[fn].get()
            if v != '':
                f.write('%s %d' % (v, 0) + '\n')
            else:
                files.remove(fn)
    # preprocess and etc
    # call gen_db
    db_name = 'data_lmdb.' + phase
    if os.path.exists(db_name):
        shutil.rmtree(db_name)
    p = subprocess.Popen(shlex.split(
        '%s/convert_imageset %s/ %s %s' % (tools_root, '.', list_file_name, db_name)
    ))
    p.wait()
    # generate label
    # +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
    # Please set the following values and paths as per your needs
    N = len(files)  # Number of data instances
    M = 4  # Number of possible labels for each data instance
    # +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
    # -------- Write in LMDB for Caffe ----------
    X = np.zeros((N, M, 1, 1), dtype=np.uint8)
    y = np.zeros(N, dtype=np.int64)
    map_size = 128 * 1024 * 1024 * 1024
    output_lmdb_path = 'label_lmdb.' + phase
    if os.path.exists(output_lmdb_path):
        shutil.rmtree(output_lmdb_path)
    env = lmdb.open(output_lmdb_path, map_size=map_size)
    for index, f in enumerate(files):
        level = file_level[f]
        for i in range(M):
            if level > i:
                X[index, i, 0, 0] = 1
            else:
                break
                # print(level, X[index,:,:,:])
    with env.begin(write=True) as txn:
        # txn is a Transaction object
        for i in range(N):
            datum = caffe.proto.caffe_pb2.Datum()
            datum.channels = X.shape[1]
            datum.height = X.shape[2]
            datum.width = X.shape[3]
            datum.data = X[i].tostring()  # or .tobytes() if numpy < 1.9
            datum.label = int(0)
            str_id = '{:08}'.format(i)
            # The encode is only essential in Python 3
            txn.put(str_id.encode('ascii'), datum.SerializeToString())
